"""01
PyTorch 线性回归实现
"""
import torch
import numpy as np

# 1.散点输入
# 1、散点输入
# 定义输入数据
data = [[-0.5, 7.7], [1.8, 98.5], [0.9, 57.8], [0.4, 39.2], [-1.4, -15.7], [-1.4, -37.3], [-1.8, -49.1], [1.5, 75.6], [0.4, 34.0], [0.8, 62.3]]
# 转换为 NumPy 数组
data = np.array(data)
# 提取 x_data 和 y_data
x_data = data[:, 0]
y_data = data[:, 1]
#将x_data 和y_data 转化成tensor
x_train=torch.tensor(x_data,dtype=torch.float32)
y_train=torch.tensor(y_data,dtype=torch.float32)

# 2.定义前向模型
#看官网https://pytorch.org/docs/stable/nn.html
#添加 torch.nn
import torch.nn as nn

#导入TensorboardX 中的SummaryWriter
from tensorboardX import SummaryWriter
#创建SummaryWriter的对象
writer=SummaryWriter(logdir="logs")
# 可视化一个神经网络




# 方案4最常用的网络结构
# 写个类继承于nn.Module
class LinearModule(nn.Module):
    def __init__(self):
        super(LinearModule,self).__init__()
        self.layers=nn.Linear(1,1)
        # self.layers1 = nn.Linear(1, 1)
    def forward(self,x):
        x=self.layers(x)
        # x = self.layers1(x)
        return x

model=LinearModule()

#3.定义损失函数和优化器
#看官网https://pytorch.org/docs/stable/nn.html
#在torch.nn下面有损失函数
criterion=nn.MSELoss()
#看官网https://pytorch.org/docs/stable/nn.html
#在torch.optim下有优化器
# 第一个参数是模型的参数 第二个是学习率 第三个动量后续学了之后再介绍
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)

#4.开始迭代
epoches =500

for n in range(1,epoches+1):
    #获得预测值
    #从1*10转成10*1
    # 看官网https://pytorch.org/docs/stable/nn.html
    # 搜 torch.unsqueeze
    # 扩维度unsqueeze
    y_hat=model(x_train.unsqueeze(1))
    #计算损失
    #需要把预测值的维度降成和输入一样的维度
    loss=criterion(y_hat.squeeze(1),y_train)
    # print(f"loss{loss}")
    #作用 清空之前存储的优化器中的梯度
    optimizer.zero_grad()
    #计算损失函数关于模型参数的梯度
    loss.backward()
    #根据优化算法去更新参数
    optimizer.step()
    # 5、显示频率设置
    #不先画图
    writer.add_scalar("loss",loss,n)
    writer.add_scalar("learing_rate", optimizer.param_groups[0]["lr"],n)
    if n % 10 == 0 or n == 1:
        print(f"epoches:{n},loss:{loss}")
        # torch.save(model.state_dict(), f"./models/model{n}.pth")

writer.add_graph(model,torch.rand(1,))
writer.close()